from tonglun_py import *
import numpy as np
import psutil
import gc

# Simple function that just returns the first argument
export_fn("identity", lambda x: x)

process = psutil.Process()
print('Initial memory:', process.memory_info().rss)

# Test the memory leak with a simple identity function
for i in range(1000):
    x = np.array([1.0, 2.0, 3.0])
    result = call_fn("identity", x)

    if i % 100 == 0:
        gc.collect()  # Force garbage collection
        print(f"Iteration {i}, memory: {process.memory_info().rss}")

gc.collect()
print('Final memory:', process.memory_info().rss) 

for i in range(1000):
    x = np.array([1.0, 2.0, 3.0])
    result = call_fn("identity", x)

    if i % 100 == 0:
        gc.collect()  # Force garbage collection
        print(f"Iteration {i}, memory: {process.memory_info().rss}")

gc.collect()
print('Final memory:', process.memory_info().rss) 
